use std::cmp::Ordering;

use reqwest::{Client, RequestBuilder};
use serde::{Deserialize, Serialize};

use crate::util::property;

#[derive(Serialize, Debug)]
struct RerankReq {
    query: String,
    documents: Vec<String>,
}

#[derive(Deserialize, Debug)]
struct RerankResp {
    model: String,
    object: String,
    usage: RerankRespUsage,
    results: Vec<RerankRespResult>,
}

#[derive(Deserialize, Debug)]
struct RerankRespUsage {
    prompt_tokens: i32,
    total_tokens: i32,
}

#[derive(Deserialize, Debug)]
struct RerankRespResult {
    index: u32,
    relevance_score: f32,
}

pub struct Reranker {
    pub end_point: String,
}

impl Reranker {
    pub async fn rerank<T: Clone>(
        &self,
        question: String,
        contents: Vec<String>,
        sort_list: Vec<T>,
    ) -> Vec<T> {
        // http客户端
        let client = Client::new();

        let req = RerankReq {
            query: question.clone(),     // todo 这里直接拷贝有性能问题
            documents: contents.clone(), // todo 这里直接拷贝有性能问题
        };

        let resp_result: Result<reqwest::Response, reqwest::Error> = client.post(&self.end_point).json(&req).send().await;
        if let Ok(resp) = resp_result {
            let rerank_scores_result: Result<RerankResp, reqwest::Error> = resp.json().await;
            if let Ok(rerank_scores) = rerank_scores_result {
                // 结果排序
                let mut rerank_results = rerank_scores.results;
                rerank_results.sort_by(|a, b| {
                    a.relevance_score
                        .partial_cmp(&b.relevance_score)
                        .unwrap_or(Ordering::Equal)
                });
                rerank_results.reverse();

                // 根据排序构建结果
                let results: Vec<T> = rerank_results
                    .iter()
                    .map(|item| item.index)
                    .map(|index| sort_list[index as usize].clone()) // todo 这里直接拷贝有性能问题
                    .collect();

                return results;
            }
        }

        // 失败，返回原始字符串
        return sort_list;
    }
}

#[test]
fn test() {
    // 阻塞运行
    let rt = tokio::runtime::Runtime::new().unwrap();

    let env: property::Env = property::Env::init();
    let reranker = Reranker {
        end_point: env.rerank_endpoint.clone(),
    };

    let sort_list = vec![
        "你的头怎么尖尖的".to_string(),
        "那我问你水草长在水里还是土里".to_string(),
        "现代人越来越注重健康饮食，倾向于选择有机食品和富含纤维的食物。这种趋势推动了市场上各类健康食品的发展，满足了消费者对高品质生活的追求。".to_string(),
        "城市中的绿化工程对于提升居民生活质量至关重要。公园、绿地不仅能美化环境，还能有效降低城市的热岛效应".to_string(),
        "具有根系的水草能够通过其根部吸收水中的多余营养物质，如氮和磷，有助于减少富营养化现象，从而改善水质。".to_string(),
        "在众多水生植物中，我们可以找到有根与无根的两种主要类型。像苦草、黑藻等就属于具有发达根系的水生植物，而某些类型的藻类则不形成真正的根。".to_string(),
        "并不是所有的水草都有根。例如，浮萍就是一种没有明显根部结构的水草类型，它们通常漂浮于水面，直接从水中吸取所需的营养物质。".to_string(),
        "根据其种类的不同，有的确实拥有根系。这些根系不仅帮助植物固定在底泥中，还能从周围环境中吸收水分和养分，支持植物生长。".to_string(),
    ];

    let future = reranker.rerank(
        "水草有没有根".to_string(),
        sort_list.clone(),
        sort_list.clone(),
    );

    let resp = rt.block_on(future);
    println!("{:?}", resp);
    assert!(
        resp.get(0).unwrap_or(&"".to_string()) == &"并不是所有的水草都有根。例如，浮萍就是一种没有明显根部结构的水草类型，它们通常漂浮于水面，直接从水中吸取所需的营养物质。".to_string()
    );
}
